{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 使用rosettaVS进行虚拟筛选的实操\n",
    "\n",
    "仅进行KLHDC2的筛选测试.\n",
    "\n",
    "涉及软件较多,数据集也较大,根据实际情况运行本笔记.\n",
    "\n",
    "## 相关链接\n",
    "\n",
    "- 论文链接:\n",
    "- rosetta github仓库\n",
    "- rosetta官网链接\n",
    "- openVS github仓库\n",
    "\n",
    "## openvs代码库结构\n",
    "\n",
    "- openvs/:openvs平台的工具模块\n",
    "- benchmarks/:casf_2016和dud两个基准集的对接筛选相关程序,目录内的readme有进一步指导\n",
    "- experiments/:两个论文介绍的虚拟筛选案例,目录内的readme有进一步指导\n",
    "- databases/:小分子库的数据示例\n",
    "- scripts/:预处理小分子库,生成指纹的脚本\n",
    "\n",
    "## 1 准备工作\n",
    "\n",
    "### 1.1 下载质心库和rosetta\n",
    "\n",
    "```bash\n",
    "git clone https://github.com/gfzhou/OpenVS.git ~/git_develop/openvs #克隆到~/git_develop/openvs下,自行调整路径\n",
    "conda create -n openvs python=3.9\n",
    "conda activate openvs\n",
    "conda install --file requirements.txt\n",
    "conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia # 建议使用镜像源\n",
    "pip install -e .\n",
    "conda install pymol-open-source ipykernel\n",
    "yay -S slurm-llnl # (可选)要尝试使用slurm调度器,用包管理器安装.注意arch系下包的名称为slurm-llnl而非slurm\n",
    "\n",
    "# 下载并编译rosetta开发版本,从略\n",
    "\n",
    "set -gx ROSETTAHOME ~/git_develop/rosetta # fish中设置为当前conda环境路径 \n",
    "export ROSETTAHOME =~/git_develop/rosetta # bash设置环境变量\n",
    "\n",
    "# conda install dimorphite-dl # 安装dimorphite-dl,conda-forge源,对openvs源代码而言不合适\n",
    "git clone https://github.com/durrantlab/dimorphite_dl.git ~/git_develop/dimorphite\n",
    "set -gx DIMORPHITE ~/git_develop/dimorphite # fish\n",
    "export DIMORPHITE =~/git_develop/dimorphite # bash\n",
    "\n",
    "# conda install --channel=https://conda.ccdc.cam.ac.uk csd-python-api # 由于该通道依赖冲突,暂时不考虑安装\n",
    "\n",
    "aria2c -s 16 -x 16 \"https://files.\n",
    "ipd.uw.edu/pub/OpenVS/centroids.tgz\" -d ~/projects # 质心库下载目录,自行调整,25GB\n",
    "\n",
    "# TODO:放置以上分子库到openvs/databases\n",
    "```\n",
    "\n",
    "由于vscode中的notebook找不到环境变量,使用os额外设置.直接运行python脚本无须额外设置:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['DIMORPHITE']=os.path.expanduser('~/git_develop/dimorphite')\n",
    "os.environ['ROSETTAHOME']=os.path.expanduser('~/git_develop/rosetta')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 下载虚拟筛选库\n",
    "\n",
    "论文使用[Enamine REAL可合成小分子库](https://enamine.net/library-synthesis/real-compounds),该网页按重原子数分割为多个库,目前共计67亿个分子.\n",
    "\n",
    "本笔记仅作演示,因此仅使用REAL的[多样性子集](https://enamine.net/compound-collections/real-compounds/real-database-subsets),含700万小分子.下载(~200M)需要注册账号并登录.\n",
    "\n",
    "下载完成后解压到下载目录(可重命名)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "smiles\tid\tMW\tHAC\tsLogP\tHBA\tHBD\tRotBonds\tFSP3\tTPSA\tlead-like\t350/3_lead-like\tfragments\tstrict_fragments\tPPI_modulators\tnatural_product-like\tType\tInChiKey\n",
      "CC(C)NC(=O)C1=CC=CO1\ts_527____153905____156108\t153.181\t11\t1.418\t2\t1\t2\t0.375\t42.240\tTrue\t\tTrue\tTrue\t\t\tS\tNJVSUHKJPNTYGL-UHFFFAOYSA-N\n",
      "COCCNC(=O)C(C)(C)Cl\tm_22____57676____19882710\t179.647\t11\t0.766\t2\t1\t4\t0.857\t38.330\tTrue\t\t\t\t\t\tM\tJPZAHQNIVFLHMM-UHFFFAOYSA-N\n",
      "COC(CN)CN1CCCC1\tm_271302____8904550____25575944\t158.245\t11\t0.056\t3\t1\t4\t1.000\t38.490\tTrue\t\t\t\t\t\tM\tMOEPNVYHZMLGIH-UHFFFAOYSA-N\n",
      "CC(C)COC(=O)CC(Cl)Cl\tm_276436____14076548____23408838\t199.077\t11\t2.379\t2\t0\t4\t0.857\t26.300\tTrue\t\t\t\t\tTrue\tM\tZQNMXDTWADSSJO-UHFFFAOYSA-N\n",
      "CSCC1=CC=C(CF)C=C1\ts_62____875850____8349548\t170.252\t11\t3.019\t1\t0\t3\t0.333\t0.000\tTrue\t\t\t\t\t\tS\tLXQJVNFOCBJMGN-UHFFFAOYSA-N\n",
      "CCN(CC)C(=O)C=CCO\tm_22____57744____13570734\t157.213\t11\t0.403\t2\t1\t4\t0.625\t40.540\tTrue\t\t\t\t\tTrue\tM\tXBNCGGGNWBNXBB-UHFFFAOYSA-N\n",
      "CC[C@H]1CC(NCCO)CN1\tm_271302____25517132____25492046\t158.245\t11\t-0.291\t3\t3\t4\t1.000\t44.290\tTrue\t\t\t\t\tTrue\tM\tHKYVXLURLQGUSW-JAMMHHFISA-N\n",
      "CCCCNC1CC(CCl)C1\tm_270004____8289202____24722536\t175.703\t11\t2.394\t1\t1\t5\t1.000\t12.030\tTrue\t\t\t\t\t\tM\tJIRDLQXLSDFJHB-UHFFFAOYSA-N\n",
      "CCOC(=O)NCCC(C)C\tm_282070____22187988____22150950\t159.229\t11\t1.779\t2\t1\t4\t0.875\t38.330\tTrue\t\t\t\t\t\tM\tPXONPKCOJADKTQ-UHFFFAOYSA-N\n"
     ]
    }
   ],
   "source": [
    "!head -n 10 ~/projects/real6.8M.cxsmiles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "为缩短测试时间,对其进行系统抽样,生成约240个分子(预计在24个CPU核心上耗费20min),并将文献中的C29加入进去:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "%mkdir -p test/inputs\n",
    "with open('/home/regen/projects/real6.8M.cxsmiles','r') as real68file,open('test/inputs/real68_240.smi','w') as real_small_file:\n",
    "    for idx,line in enumerate(real68file):\n",
    "        if idx%(6.8e6//240)==0:\n",
    "            real_small_file.write(line)\n",
    "    real_small_file.write('CC(C)(C)C=1C=CC(CSCC(=O)NCC2=CN(CC(=O)O)N=N2)=CC1 C29   \\n')# C29的smile"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 smiles质子化\n",
    "\n",
    "使用dimorphite_dl对smi文件添加质子,能够处理错误的行,所以cxsmiles的标题行可以保留."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ['PYTHONPATH']='~/git_develop/openvs' # 把openvs源码目录加入python库搜索路径中,方便调用openvs所有的脚本."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from experiments.E3L_6DO3.scripts.prepare_smiles import protonate_smiles_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[0;31mSignature:\u001b[0m \u001b[0mprotonate_smiles_dir\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msmi_dir\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutdir\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmode\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'slurm'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mDocstring:\u001b[0m <no docstring>\n",
      "\u001b[0;31mFile:\u001b[0m      ~/git_develop/openvs/experiments/E3L_6DO3/scripts/prepare_smiles.py\n",
      "\u001b[0;31mType:\u001b[0m      function"
     ]
    }
   ],
   "source": [
    "protonate_smiles_dir?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- smi_dir:包含smiles文件的目录,只有*.smi文件会被转化.\n",
    "- outdir:输出加质子的smi文件路径\n",
    "- mode:'slurm'使用slurm调度\n",
    "\n",
    "尝试slurm模式失败,缺少slurm配置文件,因此禁用slurm模式."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "rm: 无法删除 'test/outputs/*.smi': 没有那个文件或目录\n",
      "/home/regen/git_develop/dimorphite/dimorphite_dl.py:744: SyntaxWarning: \"is not\" with a literal. Did you mean \"!=\"?\n",
      "  if line is not \"\":\n",
      "WARNING: Skipping poorly formed SMILES string: smiles\tid\tMW\tHAC\tsLogP\tHBA\tHBD\tRotBonds\tFSP3\tTPSA\tlead-like\t350/3_lead-like\tfragments\tstrict_fragments\tPPI_modulators\tnatural_product-like\tType\tInChiKey\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "For help, use: python dimorphite_dl.py --help\n",
      "\n",
      "If you use Dimorphite-DL in your research, please cite:\n",
      "Ropp PJ, Kaminsky JC, Yablonski S, Durrant JD (2019) Dimorphite-DL: An\n",
      "open-source program for enumerating the ionization states of drug-like small\n",
      "molecules. J Cheminform 11:14. doi:10.1186/s13321-019-0336-9.\n",
      "\n",
      "\n",
      "PARAMETERS:\n",
      "\n",
      " label_states: False\n",
      "       max_ph: 7.4\n",
      " max_variants: 128\n",
      "       min_ph: 7.4\n",
      "  output_file: test/outputs/real68_240.prot.smi\n",
      "pka_precision: 0.1\n",
      "       silent: False\n",
      "       smiles: None\n",
      "  smiles_file: test/inputs/real68_240.smi\n",
      "         test: False\n",
      "\n"
     ]
    }
   ],
   "source": [
    "protonate_smiles_dir('test/inputs','test/outputs',mode=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import AllChem\n",
    "from functools import partial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "由于REAL多样性子集文件的结构特点,可以使用pandas处理,且分子id格式与REAL标准库不同,此处对源脚本`convert_enamine_smifn_to_fpfn`进行遮蔽:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def smi2bit(smi:str,nBits:int=1024,radius:float=2,useFeatures:bool=True, useChirality:bool=True) -> bytes:\n",
    "    '''将smi字符串转化为分子Morgan指纹向量.\n",
    "    \n",
    "    smi:smile字符串\n",
    "\n",
    "    nBits:Morgan指纹位数(1024)\n",
    "\n",
    "    radius:Morgan指纹半径(2)\n",
    "\n",
    "    useFeatures:是否使用特征不变量,生成FCFP(True),如果为False,使用原子节点不变量,生成ECFP\n",
    "\n",
    "    useChirality:是否编码smile中的手性信息(True)\n",
    "    '''\n",
    "    mol = Chem.MolFromSmiles(smi)\n",
    "    return AllChem.GetMorganFingerprintAsBitVect(mol, radius=radius, nBits=nBits, useFeatures=useFeatures, useChirality=useChirality).ToBinary()\n",
    "\n",
    "def convert_enamine_smifn_to_fpfn(smifn:str, fpfn:str, nBits:int=1024, radius:float=2, useFeature:bool=True, useChirality:bool=True, overwrite:bool=True) -> int:\n",
    "    '''对从enamine下载的多样性分子库smi文件提取smile和分子id,创建分子Morgan指纹文件.\n",
    "    \n",
    "    smifn:smi文件名\n",
    "\n",
    "    fpfn:指纹文件名\n",
    "\n",
    "    nBits:Morgan指纹位数(1024)\n",
    "\n",
    "    radius:Morgan指纹半径(2)\n",
    "\n",
    "    useFeatures:是否使用特征不变量,生成FCFP(True),如果为False,使用原子节点不变量,生成ECFP\n",
    "\n",
    "    useChirality:是否编码smile中的手性信息(True)\n",
    "\n",
    "    overwrite:是否对已经生成的指纹文件进行覆盖(True)\n",
    "    '''\n",
    "    if os.path.exists(fpfn) and not overwrite:\n",
    "        print(f\"{fpfn} exists, skip\")\n",
    "        return 0\n",
    "    \n",
    "    smi_df=pd.read_csv(smifn,usecols=['smiles','id'],delim_whitespace=True)\n",
    "    smi_df['fp_binary']=smi_df['smiles'].apply(partial(smi2bit,radius=radius, nBits=nBits, useFeatures=useFeature, useChirality=useChirality) ) # 对整列生成bit\n",
    "    smi_df.columns=['smiles', 'molecule_id', 'fp_binary'] # 重命名列,与后续流程接口保持一致\n",
    "    smi_df=smi_df.loc[:,['molecule_id','smiles', 'fp_binary']] # 调换列顺序\n",
    "    smi_df.to_feather(fpfn)\n",
    "    print(f'Saved:{fpfn}')\n",
    "    return 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved:real68_240fp.feather\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "str"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "convert_enamine_smifn_to_fpfn('real68_240.cxsmiles','real68_240fp.feather') #生成文件的fp.feather由源代码下游脚本提供"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scripts import gen_fp_enamine_real\n",
    "gen_fp_enamine_real.convert_enamine_smifn_to_fpfn = convert_enamine_smifn_to_fpfn # 运行时遮蔽enamine指纹的实现"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "script目录下gen_fp_enamine_real.convert_enamine_real_db_smifns主要用于批量生成smi文件的Morgan指纹,指定输入文件和输出文件的路径,会自行搜索路径下的子文件夹,对应生成输出文件夹.\n",
    "\n",
    "因此在openvs源码路径下,可以将下载的分子库文件移动到databases下,默认跳过已经生成指纹的文件."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[0;31mSignature:\u001b[0m\n",
      "\u001b[0mgen_fp_enamine_real\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconvert_enamine_real_db_smifns\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0minbasedir\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0moutbasedir\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0mmode\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'slurm'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mDocstring:\u001b[0m <no docstring>\n",
      "\u001b[0;31mFile:\u001b[0m      ~/git_develop/openvs/scripts/gen_fp_enamine_real.py\n",
      "\u001b[0;31mType:\u001b[0m      function"
     ]
    }
   ],
   "source": [
    "gen_fp_enamine_real.convert_enamine_real_db_smifns?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- inbasedir:输入文件的父目录\n",
    "- outbasedir:输出文件的父目录\n",
    "- mode:slurm,multiprocess,debug:用slurm,multiprocess调度并行任务,或者debug模式"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cctbx",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
